#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
股价与情感指数可视化分析工具
用于分析股票价格与舆论情感之间的关系
"""

import os
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib.ticker import MaxNLocator
import numpy as np
from datetime import datetime
import seaborn as sns
# from scipy import stats
import matplotlib as mpl
import platform
from matplotlib.font_manager import FontProperties
plt.rcParams['font.sans-serif'] = ['SimHei']
# Matplotlib中设置字体-黑体，解决Matplotlib中文乱码问题
plt.rcParams['axes.unicode_minus'] = False
# 解决Matplotlib坐标轴负号'-'显示为方块的问题
sns.set(font='SimHei')

font_path = './font/SimHei.ttf'
font_prop = FontProperties(fname=font_path)
plt.rcParams['font.family'] = font_prop.get_name()
mpl.font_manager.fontManager.addfont(font_path)

class StockSentimentVisualizer:
    """股价与情感指数可视化分析类"""
    
    def __init__(self, stock_data_path, sentiment_data_path, decay_factor=0.8):
        """
        初始化可视化器
        
        参数:
            stock_data_path: 股价数据CSV文件路径
            sentiment_data_path: 情感指数CSV文件路径
            decay_factor: 时间衰减因子，用于计算权重（0~1之间，越接近1衰减越慢）
        """
        self.stock_data_path = stock_data_path
        self.sentiment_data_path = sentiment_data_path
        self.stock_data = None
        self.sentiment_data = None
        self.merged_data = None
        self.decay_factor = decay_factor
        
        # 加载数据
        self._load_data()
        
    def _load_data(self):
        """加载股价和情感指数数据"""
        # 检查文件是否存在
        if not os.path.exists(self.stock_data_path):
            raise FileNotFoundError(f"股价数据文件不存在: {self.stock_data_path}")
        
        if not os.path.exists(self.sentiment_data_path):
            raise FileNotFoundError(f"情感指数数据文件不存在: {self.sentiment_data_path}")
        
        # 加载股价数据
        self.stock_data = pd.read_csv(self.stock_data_path)
        print(f"成功加载股价数据，共 {len(self.stock_data)} 条记录")
        
        # 加载情感指数数据
        self.sentiment_data = pd.read_csv(self.sentiment_data_path)
        print(f"成功加载情感指数数据，共 {len(self.sentiment_data)} 条记录")
        
        # 确保日期列格式正确
        self._format_date_columns()
        
        # 合并数据
        self._merge_data(self.decay_factor)
    
    def _format_date_columns(self):
        """处理日期列格式"""
        # 对股价数据处理日期列
        if "日期" in self.stock_data.columns:
            self.stock_data.rename(columns={"日期": "date"}, inplace=True)
        
        # 确保日期列格式统一
        try:
            self.stock_data['date'] = pd.to_datetime(self.stock_data['date'])
            self.sentiment_data['date'] = pd.to_datetime(self.sentiment_data['date'])
        except Exception as e:
            print(f"日期格式转换错误: {str(e)}")
            # 如果转换失败，尝试另一种方式
            try:
                self.stock_data['date'] = pd.to_datetime(self.stock_data['date'], format='%Y-%m-%d')
                self.sentiment_data['date'] = pd.to_datetime(self.sentiment_data['date'], format='%Y-%m-%d')
            except Exception as e:
                print(f"第二次尝试日期格式转换错误: {str(e)}")
    
    def _apply_time_decay_weights(self, group, decay_factor=0.8):
        """
        对同组内的情感指数应用时间衰减权重
        
        参数:
            group: 同一交易日的情感数据组
            decay_factor: 衰减因子，数值越小衰减越快
            
        返回:
            加权平均后的情感指数
        """
        if len(group) <= 1:
            return group['sentiment_index'].mean()
        
        # 将组内数据按时间排序（最近的在后面）
        group = group.sort_values(by='date')
        
        # 计算每条数据距离交易日结束的天数差
        max_date = group['date'].max()
        group['days_diff'] = (max_date - group['date']).dt.total_seconds() / (24 * 3600)
        
        # 计算指数衰减权重：越接近交易日结束，权重越大
        group['weight'] = np.power(decay_factor, group['days_diff'])
        
        # 归一化权重
        group['weight'] = group['weight'] / group['weight'].sum()
        
        # 计算加权平均
        weighted_sentiment = (group['sentiment_index'] * group['weight']).sum()
        
        return weighted_sentiment
    
    def _aggregate_by_trading_day(self, decay_factor=0.8):
        """
        按交易日聚合情感数据，并应用时间衰减权重
        
        参数:
            decay_factor: 时间衰减因子
            
        返回:
            按交易日聚合后的情感数据
        """
        # 确保情感数据有日期列
        if 'date' not in self.sentiment_data.columns:
            raise ValueError("情感指数数据缺少date列")
        
        # 提取股价数据中的交易日
        trading_days = set(self.stock_data['date'].dt.date)
        
        # 为情感数据添加一个交易日列，映射到最近的交易日
        self.sentiment_data['trading_date'] = self.sentiment_data['date'].apply(
            lambda x: self._find_nearest_trading_day(x, trading_days)
        )
        # print(self.sentiment_data['trading_date'])
        
        # 按交易日分组，应用时间衰减权重
        aggregated_sentiment = self.sentiment_data.groupby('trading_date').apply(
            lambda x: self._apply_time_decay_weights(x, decay_factor)
        ).reset_index()
        
        # 重命名列
        aggregated_sentiment.columns = ['date', 'sentiment_index']
        
        # 转换回datetime格式
        aggregated_sentiment['date'] = pd.to_datetime(aggregated_sentiment['date'])
        
        return aggregated_sentiment
    
    def _aggregate_by_trading_day_no_decay(self):
        """
        按交易日聚合情感数据，不应用时间衰减权重（简单平均）
        
        返回:
            按交易日简单平均聚合后的情感数据
        """
        # 确保情感数据有日期列
        if 'date' not in self.sentiment_data.columns:
            raise ValueError("情感指数数据缺少date列")
        
        # 提取股价数据中的交易日
        trading_days = set(self.stock_data['date'].dt.date)
        
        # 为情感数据添加一个交易日列，映射到最近的交易日
        if 'trading_date' not in self.sentiment_data.columns:
            self.sentiment_data['trading_date'] = self.sentiment_data['date'].apply(
                lambda x: self._find_nearest_trading_day(x, trading_days)
            )
        
        # 按交易日分组，计算简单平均
        aggregated_sentiment = self.sentiment_data.groupby('trading_date')['sentiment_index'].mean().reset_index()
        
        # 重命名列
        aggregated_sentiment.columns = ['date', 'sentiment_index']
        
        # 转换回datetime格式
        aggregated_sentiment['date'] = pd.to_datetime(aggregated_sentiment['date'])
        
        return aggregated_sentiment
    
    def process_data_no_decay(self):
        """
        处理数据但不应用时间衰减权重，用于对比分析
        
        返回:
            处理后的数据
        """
        # 确保两个数据集都有日期列
        if 'date' not in self.stock_data.columns:
            raise ValueError("股价数据缺少date列")
        
        # 处理股价数据中的列名
        if "收盘" in self.stock_data.columns and "close_price" not in self.stock_data.columns:
            self.stock_data.rename(columns={"收盘": "close_price"}, inplace=True)
        elif "Close" in self.stock_data.columns and "close_price" not in self.stock_data.columns:
            self.stock_data.rename(columns={"Close": "close_price"}, inplace=True)
        
        # 按交易日聚合情感数据（无时间衰减权重）
        aggregated_sentiment = self._aggregate_by_trading_day_no_decay()
        
        # 使用日期列合并数据
        merged_data = pd.merge(
            self.stock_data[['date', 'close_price']], 
            aggregated_sentiment, 
            on='date', 
            how='inner'
        )
        
        # 将日期设置为索引
        merged_data.set_index('date', inplace=True)
        
        # 按日期排序
        merged_data.sort_index(inplace=True)
        merged_data = merged_data[-125:]
        
        # 添加归一化处理的数据列
        
        # 1. Min-Max归一化 (0-1区间)
        merged_data['close_price_norm_minmax'] = (merged_data['close_price'] - merged_data['close_price'].min()) / (merged_data['close_price'].max() - merged_data['close_price'].min())
        merged_data['sentiment_norm_minmax'] = (merged_data['sentiment_index'] - merged_data['sentiment_index'].min()) / (merged_data['sentiment_index'].max() - merged_data['sentiment_index'].min())
        
        # 2. Z-score标准化 (均值0，标准差1)
        merged_data['close_price_norm_zscore'] = (merged_data['close_price'] - merged_data['close_price'].mean()) / merged_data['close_price'].std()
        merged_data['sentiment_norm_zscore'] = (merged_data['sentiment_index'] - merged_data['sentiment_index'].mean()) / merged_data['sentiment_index'].std()
        
        # 3. 相对变化率
        merged_data['close_price_pct_change'] = merged_data['close_price'].pct_change() * 100
        merged_data['sentiment_pct_change'] = merged_data['sentiment_index'].pct_change() * 100
        
        return merged_data
    
    def _find_nearest_trading_day(self, date, trading_days):
        """
        找到最近的交易日
        
        参数:
            date: 日期
            trading_days: 交易日集合
            
        返回:
            最近的交易日
        """
        date_obj = date.date()
        
        # 如果当前日期是交易日，直接返回
        if date_obj in trading_days:
            return date_obj
        
        # 查找之前的交易日（非交易时段发帖归属到之前的交易日）
        prev_days = [day for day in trading_days if day < date_obj]
        if prev_days:
            return max(prev_days)
        
        # 如果没有之前的交易日，则归属到之后最近的交易日
        next_days = [day for day in trading_days if day > date_obj]
        if next_days:
            return min(next_days)
        
        # 如果都没有，返回原日期
        return date_obj
    
    def _merge_data(self, decay_factor=0.8):
        """
        合并股价和情感指数数据，应用交易日聚合和时间衰减权重
        
        参数:
            decay_factor: 时间衰减因子
        """
        # 确保两个数据集都有日期列
        if 'date' not in self.stock_data.columns:
            raise ValueError("股价数据缺少date列")
        
        # 处理股价数据中的列名
        if "收盘" in self.stock_data.columns:
            self.stock_data.rename(columns={"收盘": "close_price"}, inplace=True)
        elif "Close" in self.stock_data.columns:
            self.stock_data.rename(columns={"Close": "close_price"}, inplace=True)
        
        # 按交易日聚合情感数据，并应用时间衰减权重
        aggregated_sentiment = self._aggregate_by_trading_day(decay_factor)
        
        # 使用日期列合并数据
        self.merged_data = pd.merge(
            self.stock_data[['date', 'close_price']], 
            aggregated_sentiment, 
            on='date', 
            how='inner'
        )
        
        # 将日期设置为索引
        self.merged_data.set_index('date', inplace=True)
        
        # 按日期排序
        self.merged_data.sort_index(inplace=True)
        self.merged_data = self.merged_data[-30:]
        
        # 添加归一化处理的数据列
        self._add_normalized_columns()
        
        print(f"数据合并完成，共 {len(self.merged_data)} 条记录")
        print(f"应用了交易日聚合和时间衰减权重（衰减因子={decay_factor}）")
        print(f"已添加归一化处理的数据列")
        print(self.merged_data.head())
    
    def _add_normalized_columns(self):
        """添加归一化处理的数据列"""
        if self.merged_data is None or len(self.merged_data) == 0:
            return
        
        # 1. Min-Max归一化 (0-1区间)
        self.merged_data['close_price_norm_minmax'] = (self.merged_data['close_price'] - self.merged_data['close_price'].min()) / (self.merged_data['close_price'].max() - self.merged_data['close_price'].min())
        self.merged_data['sentiment_norm_minmax'] = (self.merged_data['sentiment_index'] - self.merged_data['sentiment_index'].min()) / (self.merged_data['sentiment_index'].max() - self.merged_data['sentiment_index'].min())
        
        # 2. Z-score标准化 (均值0，标准差1)
        self.merged_data['close_price_norm_zscore'] = (self.merged_data['close_price'] - self.merged_data['close_price'].mean()) / self.merged_data['close_price'].std()
        self.merged_data['sentiment_norm_zscore'] = (self.merged_data['sentiment_index'] - self.merged_data['sentiment_index'].mean()) / self.merged_data['sentiment_index'].std()
        
        # 3. 相对变化率
        self.merged_data['close_price_pct_change'] = self.merged_data['close_price'].pct_change() * 100
        self.merged_data['sentiment_pct_change'] = self.merged_data['sentiment_index'].pct_change() * 100
    
    def plot_trend(self, title="股价与情感指数走势图", figsize=(14, 8), save_path=None, dpi=100):
        """
        绘制股价和情感指数的走势图
        
        参数:
            title: 图表标题
            figsize: 图表大小
            save_path: 保存路径，如果为None则不保存
            dpi: 图像分辨率
        """
        if self.merged_data is None or len(self.merged_data) == 0:
            print("没有可用的数据进行绘图")
            return
        
        # 创建图表
        fig, ax1 = plt.subplots(figsize=figsize)
        
        # 设置主题样式
        # plt.style.use('seaborn-v0_8-whitegrid')
        
        # 绘制股价曲线
        color = 'tab:red'
        ax1.set_xlabel('日期')
        ax1.set_ylabel('股价', color=color)
        ax1.plot(self.merged_data.index, self.merged_data['close_price'], color=color, linewidth=2, marker='o', label='股价')
        ax1.tick_params(axis='y', labelcolor=color)
        
        # 设置x轴日期格式
        ax1.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
        plt.xticks(rotation=45)
        
        # 创建第二个y轴
        ax2 = ax1.twinx()
        color = 'tab:blue'
        ax2.set_ylabel('情感指数', color=color)
        ax2.plot(self.merged_data.index, self.merged_data['sentiment_index'], color=color, linewidth=2, marker='x', label='情感指数')
        ax2.tick_params(axis='y', labelcolor=color)
        
        # 添加网格线
        ax1.grid(True, linestyle='--', alpha=0.6)
        
        # 添加图例
        lines1, labels1 = ax1.get_legend_handles_labels()
        lines2, labels2 = ax2.get_legend_handles_labels()
        ax1.legend(lines1 + lines2, labels1 + labels2, loc='upper left')
        
        # 设置标题
        plt.title(title, fontsize=15)
        
        # 调整布局
        plt.tight_layout()
        
        # 保存图表
        if save_path:
            plt.savefig(save_path, dpi=dpi, bbox_inches='tight')
            print(f"图表已保存至: {save_path}")
        
        # 显示图表
        plt.show()
    
    def plot_correlation(self, figsize=(14, 10), save_path=None, dpi=100):
        """
        绘制股价和情感指数的相关性分析图
        
        参数:
            figsize: 图表大小
            save_path: 保存路径，如果为None则不保存
            dpi: 图像分辨率
        """
        if self.merged_data is None or len(self.merged_data) < 5:
            print("数据不足，无法进行相关性分析")
            return
        
        # 创建图表
        fig, axes = plt.subplots(2, 2, figsize=figsize)
        
        # 设置主题样式
        # plt.style.use('seaborn-v0_8-whitegrid')
        
        # 1. 计算相关系数
        correlation = self.merged_data['close_price'].corr(self.merged_data['sentiment_index'])
        corr_text = f"相关系数: {correlation:.4f}"
        
        # 2. 散点图和回归线
        sns.regplot(
            x='sentiment_index', 
            y='close_price', 
            data=self.merged_data, 
            ax=axes[0, 0], 
            scatter_kws={'alpha':0.6, 's':50},
            line_kws={'color': 'red'}
        )
        axes[0, 0].set_title(f'股价与情感指数散点图 ({corr_text})')
        axes[0, 0].set_xlabel('情感指数')
        axes[0, 0].set_ylabel('股价')
        
        # 3. 热力图
        sns.heatmap(
            self.merged_data[['close_price', 'sentiment_index']].corr(),
            annot=True,
            cmap='coolwarm',
            ax=axes[0, 1],
            vmin=-1, vmax=1
        )
        axes[0, 1].set_title('相关性热力图')
        
        # 4. 股价变化与情感指数的关系
        # 计算每日股价变化百分比
        self.merged_data['price_change_pct'] = self.merged_data['close_price'].pct_change() * 100
        
        # 去除NaN值
        data_change = self.merged_data.dropna()
        
        # 散点图和回归线
        change_corr = data_change['price_change_pct'].corr(data_change['sentiment_index'])
        change_corr_text = f"相关系数: {change_corr:.4f}"
        
        sns.regplot(
            x='sentiment_index', 
            y='price_change_pct', 
            data=data_change, 
            ax=axes[1, 0], 
            scatter_kws={'alpha':0.6, 's':50},
            line_kws={'color': 'red'}
        )
        axes[1, 0].set_title(f'股价变化百分比与情感指数 ({change_corr_text})')
        axes[1, 0].set_xlabel('情感指数')
        axes[1, 0].set_ylabel('股价变化百分比 (%)')
        
        # 5. 情感指数分布直方图
        sns.histplot(self.merged_data['sentiment_index'], kde=True, ax=axes[1, 1])
        axes[1, 1].set_title('情感指数分布')
        axes[1, 1].set_xlabel('情感指数')
        axes[1, 1].set_ylabel('频率')
        
        # 调整布局
        plt.tight_layout()
        
        # 保存图表
        if save_path:
            plt.savefig(save_path, dpi=dpi, bbox_inches='tight')
            print(f"相关性分析图已保存至: {save_path}")
        
        # 显示图表
        plt.show()
    
    def lag_analysis(self, max_lag=10, figsize=(14, 6), save_path=None, dpi=100):
        """
        滞后相关性分析
        
        参数:
            max_lag: 最大滞后天数
            figsize: 图表大小
            save_path: 保存路径，如果为None则不保存
            dpi: 图像分辨率
        """
        if self.merged_data is None or len(self.merged_data) <= max_lag:
            print("数据不足，无法进行滞后分析")
            return
        
        # 计算不同滞后期的相关系数
        correlations = []
        for i in range(max_lag + 1):
            if i == 0:
                # 当前相关性
                corr = self.merged_data['close_price'].corr(self.merged_data['sentiment_index'])
            else:
                # 滞后相关性 (股价滞后于情感)
                corr = self.merged_data['close_price'].corr(self.merged_data['sentiment_index'].shift(-i))
            correlations.append(corr)
        
        # 创建图表
        plt.figure(figsize=figsize)

        # plt.style.use('seaborn-v0_8-whitegrid')
        
        plt.bar(range(max_lag + 1), correlations, color='skyblue', alpha=0.7)
        plt.plot(range(max_lag + 1), correlations, 'ro-', linewidth=2)
        
        plt.xlabel('滞后天数')
        plt.ylabel('相关系数')
        plt.title('股价与情感指数的滞后相关性分析')
        plt.xticks(range(max_lag + 1))
        plt.axhline(y=0, color='k', linestyle='--', alpha=0.3)
        
        # 添加数值标签
        for i, v in enumerate(correlations):
            plt.text(i, v + 0.02 * (1 if v >= 0 else -1), f'{v:.3f}', ha='center')
        
        plt.grid(True, linestyle='--', alpha=0.6)
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=dpi, bbox_inches='tight')
            print(f"滞后相关性分析图已保存至: {save_path}")
        
        plt.show()
        
        # 找出最大相关系数对应的滞后期
        max_corr_idx = np.argmax(np.abs(correlations))
        max_corr = correlations[max_corr_idx]
        print(f"最强相关性在滞后 {max_corr_idx} 天时出现，相关系数为 {max_corr:.4f}")
        
        return correlations
    
    def lag_analysis_normalized(self, normalization_type='zscore', max_lag=10, figsize=(14, 6), save_path=None, dpi=100):
        """
        归一化数据的滞后相关性分析
        
        参数:
            normalization_type: 归一化类型，可选'minmax'、'zscore'或'pct_change'
            max_lag: 最大滞后天数
            figsize: 图表大小
            save_path: 保存路径，如果为None则不保存
            dpi: 图像分辨率
        """
        if self.merged_data is None or len(self.merged_data) <= max_lag:
            print("数据不足，无法进行滞后分析")
            return
        
        # 选择对应的归一化数据列
        if normalization_type == 'minmax':
            price_col = 'close_price_norm_minmax'
            sentiment_col = 'sentiment_norm_minmax'
            type_desc = "MinMax归一化"
        elif normalization_type == 'zscore':
            price_col = 'close_price_norm_zscore'
            sentiment_col = 'sentiment_norm_zscore'
            type_desc = "Z-score标准化"
        elif normalization_type == 'pct_change':
            price_col = 'close_price_pct_change'
            sentiment_col = 'sentiment_pct_change'
            type_desc = "相对变化率"
        else:
            print(f"不支持的归一化类型: {normalization_type}")
            return
            
        # 准备数据，去除NaN值
        data = self.merged_data.dropna(subset=[price_col, sentiment_col])
        
        # 计算不同滞后期的相关系数
        correlations = []
        for i in range(max_lag + 1):
            if i == 0:
                corr = data[price_col].corr(data[sentiment_col])
            else:
                corr = data[price_col].corr(data[sentiment_col].shift(-i))
            correlations.append(corr)
        
        # 创建图表
        plt.figure(figsize=figsize)
        
        plt.bar(range(max_lag + 1), correlations, color='skyblue', alpha=0.7)
        plt.plot(range(max_lag + 1), correlations, 'ro-', linewidth=2)
        
        plt.xlabel('滞后天数')
        plt.ylabel('相关系数')
        plt.title(f'股价与情感指数的滞后相关性分析 ({type_desc})')
        plt.xticks(range(max_lag + 1))
        plt.axhline(y=0, color='k', linestyle='--', alpha=0.3)
        
        # 添加数值标签
        for i, v in enumerate(correlations):
            plt.text(i, v + 0.02 * (1 if v >= 0 else -1), f'{v:.3f}', ha='center')
        
        plt.grid(True, linestyle='--', alpha=0.6)
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=dpi, bbox_inches='tight')
            print(f"归一化滞后相关性分析图已保存至: {save_path}")
        
        plt.show()
        
        # 找出最大相关系数对应的滞后期
        max_corr_idx = np.argmax(np.abs(correlations))
        max_corr = correlations[max_corr_idx]
        print(f"[{type_desc}] 最强相关性在滞后 {max_corr_idx} 天时出现，相关系数为 {max_corr:.4f}")
        
        return correlations
    
    def compare_normalizations_lag_effect(self, max_lag=5, figsize=(14, 10), save_path=None, dpi=100):
        """
        比较不同归一化方法下的滞后效应
        
        参数:
            max_lag: 最大滞后天数
            figsize: 图表大小
            save_path: 保存路径，如果为None则不保存
            dpi: 图像分辨率
        """
        if self.merged_data is None or len(self.merged_data) <= max_lag:
            print("数据不足，无法进行滞后分析比较")
            return
        
        # 定义归一化方法
        norm_types = [
            ('原始数据', 'close_price', 'sentiment_index'),
            ('MinMax归一化', 'close_price_norm_minmax', 'sentiment_norm_minmax'),
            ('Z-score标准化', 'close_price_norm_zscore', 'sentiment_norm_zscore'),
            ('相对变化率', 'close_price_pct_change', 'sentiment_pct_change')
        ]
        
        # 创建图表
        fig, axes = plt.subplots(2, 2, figsize=figsize)
        axes = axes.flatten()
        
        # 对每种归一化方法计算滞后相关性
        for idx, (label, price_col, sentiment_col) in enumerate(norm_types):
            # 准备数据，去除NaN值
            data = self.merged_data.dropna(subset=[price_col, sentiment_col])
            
            # 计算不同滞后期的相关系数
            correlations = []
            for i in range(max_lag + 1):
                if i == 0:
                    corr = data[price_col].corr(data[sentiment_col])
                else:
                    corr = data[price_col].corr(data[sentiment_col].shift(-i))
                correlations.append(corr)
            
            # 绘制到对应的子图
            ax = axes[idx]
            ax.bar(range(max_lag + 1), correlations, color='skyblue', alpha=0.7)
            ax.plot(range(max_lag + 1), correlations, 'ro-', linewidth=2)
            
            ax.set_xlabel('滞后天数')
            ax.set_ylabel('相关系数')
            ax.set_title(f'{label}')
            ax.set_xticks(range(max_lag + 1))
            ax.axhline(y=0, color='k', linestyle='--', alpha=0.3)
            
            # 添加数值标签
            for i, v in enumerate(correlations):
                ax.text(i, v + 0.02 * (1 if v >= 0 else -1), f'{v:.3f}', ha='center')
            
            ax.grid(True, linestyle='--', alpha=0.6)
            
            # 找出最大相关系数
            max_corr_idx = np.argmax(np.abs(correlations))
            max_corr = correlations[max_corr_idx]
            ax.annotate(f"最大: 滞后{max_corr_idx}天, 相关系数={max_corr:.3f}", 
                      xy=(0.05, 0.05), xycoords='axes fraction', fontsize=9,
                      bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="gray", alpha=0.8))
        
        plt.tight_layout()
        
        if save_path:
            plt.savefig(save_path, dpi=dpi, bbox_inches='tight')
            print(f"归一化方法滞后效应比较图已保存至: {save_path}")
        
        plt.show()
    
    def add_smoothed_sentiment(self):
        """添加平滑处理后的情感指数"""
        windows = [3, 7, 14]  # 不同时间窗口
        for window in windows:
            self.merged_data[f'sentiment_ma{window}'] = self.merged_data['sentiment_index'].rolling(window=window).mean()
        
        # 也可以尝试指数移动平均
        self.merged_data['sentiment_ema7'] = self.merged_data['sentiment_index'].ewm(span=7).mean()

    def time_scale_analysis(self, save_path=None):
        """多尺度时间分析"""
        # 按周/月聚合数据
        weekly_data = self.merged_data.resample('W').mean()
        monthly_data = self.merged_data.resample('ME').mean()  # 使用'ME'代替'M'
        
        # 绘制不同时间尺度下的关系图
        fig, axes = plt.subplots(3, 1, figsize=(14, 18))
        
        # 日尺度
        self._plot_dual_axis(self.merged_data, axes[0], '日频数据')
        # 周尺度
        self._plot_dual_axis(weekly_data, axes[1], '周频数据')
        # 月尺度
        self._plot_dual_axis(monthly_data, axes[2], '月频数据')
        
        plt.tight_layout()
        if save_path:
            plt.savefig(save_path)
        plt.show()

    def cross_correlation_analysis(self, max_lag=30):
        """交叉相关分析，找出最佳领先/滞后关系"""
        from statsmodels.tsa.stattools import ccf
        
        # 标准化数据
        price_norm = (self.merged_data['close_price'] - self.merged_data['close_price'].mean()) / self.merged_data['close_price'].std()
        sentiment_norm = (self.merged_data['sentiment_index'] - self.merged_data['sentiment_index'].mean()) / self.merged_data['sentiment_index'].std()
        
        # 计算交叉相关
        ccf_result = ccf(price_norm.dropna(), sentiment_norm.dropna(), adjusted=False)
        
        # 确保我们只使用有效的CCF结果范围
        valid_lags = min(max_lag, len(ccf_result))
        half_lags = valid_lags // 2
        
        # 获取实际有效的CCF结果
        ccf_values = ccf_result[:valid_lags]
        
        # 创建对应的横坐标
        x_range = list(range(-half_lags, valid_lags - half_lags))
        
        # 确保x和y长度相同
        if len(x_range) > len(ccf_values):
            x_range = x_range[:len(ccf_values)]
        elif len(ccf_values) > len(x_range):
            ccf_values = ccf_values[:len(x_range)]
        
        # 找出最大相关及其滞后期
        max_corr_idx = np.argmax(np.abs(ccf_values))
        lag = x_range[max_corr_idx]
        
        plt.figure(figsize=(12, 6))
        plt.bar(x_range, ccf_values, alpha=0.7)
        plt.axvline(x=lag, color='r', linestyle='--')
        plt.title(f'股价与情感指数交叉相关分析 (最大相关滞后期: {lag}天)')
        plt.xlabel('滞后天数 (负值表示情感指数领先股价)')
        plt.ylabel('相关系数')
        plt.grid(True, alpha=0.3)
        plt.show()

    def segment_analysis(self, n_segments=3):
        """将时间序列分成多个区段进行分析"""
        # 将数据分成n_segments个时间段
        dates = self.merged_data.index
        segment_size = len(dates) // n_segments
        
        plt.figure(figsize=(15, 10))
        
        for i in range(n_segments):
            start_idx = i * segment_size
            end_idx = (i+1) * segment_size if i < n_segments-1 else len(dates)
            
            segment_data = self.merged_data.iloc[start_idx:end_idx]
            corr = segment_data['close_price'].corr(segment_data['sentiment_index'])
            
            plt.subplot(n_segments, 1, i+1)
            plt.plot(segment_data.index, segment_data['close_price'], 'r-', label='股价')
            plt.title(f'时间段 {i+1}: {segment_data.index[0].date()} 至 {segment_data.index[-1].date()}, 相关系数={corr:.4f}')
            
            plt2 = plt.twinx()
            plt2.plot(segment_data.index, segment_data['sentiment_index'], 'b-', label='情感指数')
            
            # 合并图例
            lines1, labels1 = plt.gca().get_legend_handles_labels()
            lines2, labels2 = plt2.get_legend_handles_labels()
            plt2.legend(lines1 + lines2, labels1 + labels2, loc='upper right')
        
        plt.tight_layout()
        plt.show()

    def granger_causality_test(self, max_lag=10):
        """格兰杰因果检验，分析情感指数是否对股价有预测作用"""
        from statsmodels.tsa.stattools import grangercausalitytests
        
        # 准备数据
        data = pd.DataFrame({
            'price': self.merged_data['close_price'],
            'sentiment': self.merged_data['sentiment_index']
        })
        data = data.dropna()
        
        # 情感指数是否对股价有格兰杰因果关系
        print("检验情感指数是否'引起'股价变动:")
        grangercausalitytests(data[['price', 'sentiment']], max_lag, verbose=True)
        
        # 股价是否对情感指数有格兰杰因果关系
        print("\n检验股价是否'引起'情感指数变动:")
        grangercausalitytests(data[['sentiment', 'price']], max_lag, verbose=True)

    def volatility_correlation(self, window=7):
        """分析股价波动率与情感指数波动性的关系"""
        # 计算股价的波动率(用简单的滚动标准差)
        self.merged_data['price_volatility'] = self.merged_data['close_price'].rolling(window).std()
        
        # 计算情感指数的波动性
        self.merged_data['sentiment_volatility'] = self.merged_data['sentiment_index'].rolling(window).std()
        
        # 去除NaN
        vol_data = self.merged_data.dropna()
        
        # 计算相关系数
        corr = vol_data['price_volatility'].corr(vol_data['sentiment_volatility'])
        
        # 绘制散点图
        plt.figure(figsize=(10, 8))
        plt.scatter(vol_data['sentiment_volatility'], vol_data['price_volatility'], alpha=0.6)
        plt.title(f'股价波动率与情感指数波动性相关分析 (相关系数: {corr:.4f})')
        plt.xlabel('情感指数波动性')
        plt.ylabel('股价波动率')
        
        # 添加回归线
        z = np.polyfit(vol_data['sentiment_volatility'], vol_data['price_volatility'], 1)
        p = np.poly1d(z)
        plt.plot(vol_data['sentiment_volatility'], p(vol_data['sentiment_volatility']), "r--")
        
        plt.grid(True, alpha=0.3)
        plt.show()

    def extreme_sentiment_event_analysis(self, threshold=30):
        """极端情感事件对股价的影响分析"""
        # 找出情感指数极高和极低的日期
        high_sentiment = self.merged_data[self.merged_data['sentiment_index'] > threshold].index
        low_sentiment = self.merged_data[self.merged_data['sentiment_index'] < -threshold].index
        
        # 分析这些日期后n天的股价表现
        windows = [1, 3, 5, 10]
        
        results = pd.DataFrame(index=['极高情感后', '极低情感后'], columns=[f'{w}天后平均涨跌幅(%)' for w in windows])
        
        for window in windows:
            high_returns = []
            low_returns = []
            
            for date in high_sentiment:
                try:
                    future_date = self.merged_data.index[self.merged_data.index.get_loc(date) + window]
                    pct_change = (self.merged_data.loc[future_date, 'close_price'] / 
                                 self.merged_data.loc[date, 'close_price'] - 1) * 100
                    high_returns.append(pct_change)
                except:
                    pass
                
            for date in low_sentiment:
                try:
                    future_date = self.merged_data.index[self.merged_data.index.get_loc(date) + window]
                    pct_change = (self.merged_data.loc[future_date, 'close_price'] / 
                                 self.merged_data.loc[date, 'close_price'] - 1) * 100
                    low_returns.append(pct_change)
                except:
                    pass
            
            results.loc['极高情感后', f'{window}天后平均涨跌幅(%)'] = np.mean(high_returns) if high_returns else np.nan
            results.loc['极低情感后', f'{window}天后平均涨跌幅(%)'] = np.mean(low_returns) if low_returns else np.nan
        
        return results

    def visualize_time_decay(self, days=10, decay_factor=None, figsize=(12, 6), save_path=None):
        """
        可视化时间衰减权重效果
        
        参数:
            days: 显示的天数
            decay_factor: 衰减因子，如果为None则使用实例初始化时的因子
            figsize: 图表大小
            save_path: 保存路径
        """
        if decay_factor is None:
            decay_factor = self.decay_factor
        
        # 创建示例时间序列
        dates = pd.date_range(end=pd.Timestamp.now().date(), periods=days)
        days_diff = [(dates[-1] - date).days for date in dates]
        
        # 计算权重
        weights = np.power(decay_factor, days_diff)
        weights = weights / weights.sum()  # 归一化
        
        # 创建图表
        plt.figure(figsize=figsize)
        plt.bar(range(days), weights, alpha=0.7)
        plt.plot(range(days), weights, 'ro-', linewidth=2)
        
        # 添加权重标签
        for i, w in enumerate(weights):
            plt.text(i, w + 0.01, f'{w:.3f}', ha='center')
        
        plt.xlabel('距离交易日的天数')
        plt.ylabel('权重')
        plt.title(f'时间衰减权重可视化 (衰减因子={decay_factor})')
        plt.xticks(range(days), [f'T-{d}' for d in days_diff])
        plt.grid(True, linestyle='--', alpha=0.6)
        
        if save_path:
            plt.savefig(save_path, dpi=100, bbox_inches='tight')
            print(f"时间衰减权重图已保存至: {save_path}")
        
        plt.show()

    def _plot_dual_axis(self, data, ax, title):
        """
        在指定的轴上绘制股价和情感指数的双轴图
        
        参数:
            data: 包含股价和情感指数的DataFrame
            ax: 要绘制的matplotlib轴对象
            title: 图表标题
        """
        # 绘制股价曲线
        color1 = 'tab:red'
        ax.set_xlabel('日期')
        ax.set_ylabel('股价', color=color1)
        ax.plot(data.index, data['close_price'], color=color1, linewidth=2, marker='o', label='股价')
        ax.tick_params(axis='y', labelcolor=color1)
        
        # 设置x轴日期格式
        ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
        plt.setp(ax.xaxis.get_majorticklabels(), rotation=45)
        
        # 创建第二个y轴
        ax2 = ax.twinx()
        color2 = 'tab:blue'
        ax2.set_ylabel('情感指数', color=color2)
        ax2.plot(data.index, data['sentiment_index'], color=color2, linewidth=2, marker='x', label='情感指数')
        ax2.tick_params(axis='y', labelcolor=color2)
        
        # 添加网格线
        ax.grid(True, linestyle='--', alpha=0.6)
        
        # 添加图例
        lines1, labels1 = ax.get_legend_handles_labels()
        lines2, labels2 = ax2.get_legend_handles_labels()
        ax.legend(lines1 + lines2, labels1 + labels2, loc='upper left')
        
        # 设置标题
        ax.set_title(title)
        
        # 计算相关系数
        corr = data['close_price'].corr(data['sentiment_index'])
        ax.annotate(f"相关系数: {corr:.4f}", xy=(0.02, 0.95), xycoords='axes fraction', 
                   fontsize=10, bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="gray", alpha=0.8))

    def plot_trend_normalized(self, normalization_type='zscore', title=None, figsize=(14, 8), save_path=None, dpi=100):
        """
        绘制归一化后的股价和情感指数走势图
        
        参数:
            normalization_type: 归一化类型，可选'minmax'、'zscore'或'pct_change'
            title: 图表标题
            figsize: 图表大小
            save_path: 保存路径，如果为None则不保存
            dpi: 图像分辨率
        """
        if self.merged_data is None or len(self.merged_data) == 0:
            print("没有可用的数据进行绘图")
            return
        
        # 选择对应的归一化数据列
        if normalization_type == 'minmax':
            price_col = 'close_price_norm_minmax'
            sentiment_col = 'sentiment_norm_minmax'
            type_desc = "MinMax归一化"
            ylabel = "归一化值 (0-1)"
        elif normalization_type == 'zscore':
            price_col = 'close_price_norm_zscore'
            sentiment_col = 'sentiment_norm_zscore'
            type_desc = "Z-score标准化"
            ylabel = "标准化值 (z-score)"
        elif normalization_type == 'pct_change':
            price_col = 'close_price_pct_change'
            sentiment_col = 'sentiment_pct_change'
            type_desc = "相对变化率"
            ylabel = "变化率 (%)"
        else:
            print(f"不支持的归一化类型: {normalization_type}")
            return
        
        # 过滤掉NaN值
        plot_data = self.merged_data.dropna(subset=[price_col, sentiment_col])
        
        # 如果没有提供标题，生成默认标题
        if title is None:
            title = f"股价与情感指数走势对比图 ({type_desc})"
        
        # 创建图表
        plt.figure(figsize=figsize)
        
        # 绘制归一化后的曲线
        plt.plot(plot_data.index, plot_data[price_col], 'r-', label='股价', linewidth=2)
        plt.plot(plot_data.index, plot_data[sentiment_col], 'b-', label='情感指数', linewidth=2)
        
        # 设置x轴日期格式
        plt.gca().xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
        plt.xticks(rotation=45)
        
        # 添加网格线和图例
        plt.grid(True, linestyle='--', alpha=0.6)
        plt.legend(loc='best')
        
        # 设置标题和标签
        plt.title(title, fontsize=15)
        plt.xlabel('日期')
        plt.ylabel(ylabel)
        
        # 添加相关系数注释
        corr = plot_data[price_col].corr(plot_data[sentiment_col])
        plt.annotate(f"相关系数: {corr:.4f}", xy=(0.02, 0.95), xycoords='axes fraction',
                   fontsize=12, bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="gray", alpha=0.8))
        
        # 调整布局
        plt.tight_layout()
        
        # 保存图表
        if save_path:
            plt.savefig(save_path, dpi=dpi, bbox_inches='tight')
            print(f"归一化走势图已保存至: {save_path}")
        
        # 显示图表
        plt.show()
        
        return corr
    
    def plot_correlation_normalized(self, figsize=(14, 10), save_path=None, dpi=100):
        """
        绘制归一化后的股价和情感指数相关性分析图
        
        参数:
            figsize: 图表大小
            save_path: 保存路径，如果为None则不保存
            dpi: 图像分辨率
        """
        if self.merged_data is None or len(self.merged_data) < 5:
            print("数据不足，无法进行相关性分析")
            return
        
        # 创建图表
        fig, axes = plt.subplots(2, 2, figsize=figsize)
        
        # 准备数据
        data = self.merged_data.dropna()
        
        # 1. MinMax归一化的散点图
        sns.regplot(
            x='sentiment_norm_minmax', 
            y='close_price_norm_minmax', 
            data=data, 
            ax=axes[0, 0], 
            scatter_kws={'alpha':0.6, 's':50},
            line_kws={'color': 'red'}
        )
        corr1 = data['close_price_norm_minmax'].corr(data['sentiment_norm_minmax'])
        axes[0, 0].set_title(f'MinMax归一化 (相关系数: {corr1:.4f})')
        axes[0, 0].set_xlabel('情感指数 (归一化)')
        axes[0, 0].set_ylabel('股价 (归一化)')
        
        # 2. Z-score标准化的散点图
        sns.regplot(
            x='sentiment_norm_zscore', 
            y='close_price_norm_zscore', 
            data=data, 
            ax=axes[0, 1], 
            scatter_kws={'alpha':0.6, 's':50},
            line_kws={'color': 'red'}
        )
        corr2 = data['close_price_norm_zscore'].corr(data['sentiment_norm_zscore'])
        axes[0, 1].set_title(f'Z-score标准化 (相关系数: {corr2:.4f})')
        axes[0, 1].set_xlabel('情感指数 (标准化)')
        axes[0, 1].set_ylabel('股价 (标准化)')
        
        # 3. 变化率的散点图
        sns.regplot(
            x='sentiment_pct_change', 
            y='close_price_pct_change', 
            data=data.dropna(), 
            ax=axes[1, 0], 
            scatter_kws={'alpha':0.6, 's':50},
            line_kws={'color': 'red'}
        )
        corr3 = data['close_price_pct_change'].corr(data['sentiment_pct_change'])
        axes[1, 0].set_title(f'相对变化率 (相关系数: {corr3:.4f})')
        axes[1, 0].set_xlabel('情感指数变化率 (%)')
        axes[1, 0].set_ylabel('股价变化率 (%)')
        
        # 4. 热力图
        corr_matrix = data[[
            'close_price', 'sentiment_index',
            'close_price_norm_minmax', 'sentiment_norm_minmax',
            'close_price_norm_zscore', 'sentiment_norm_zscore',
            'close_price_pct_change', 'sentiment_pct_change'
        ]].corr()
        
        sns.heatmap(
            corr_matrix,
            annot=True,
            cmap='coolwarm',
            ax=axes[1, 1],
            vmin=-1, vmax=1,
            annot_kws={"size": 8},
            fmt='.2f'
        )
        axes[1, 1].set_title('各指标相关性热力图')
        
        # 调整标签字体大小
        for i in range(2):
            for j in range(2):
                axes[i, j].tick_params(axis='both', which='major', labelsize=9)
        
        # 调整布局
        plt.tight_layout()
        
        # 保存图表
        if save_path:
            plt.savefig(save_path, dpi=dpi, bbox_inches='tight')
            print(f"归一化相关性分析图已保存至: {save_path}")
        
        # 显示图表
        plt.show()
        
        return {
            'minmax': corr1,
            'zscore': corr2,
            'pct_change': corr3
        }

    def compare_decay_vs_no_decay(self, normalization_type='zscore', figsize=(16, 12), save_path=None, dpi=100):
        """
        比较时间衰减权重与无衰减权重的情感指数对股价预测能力
        
        参数:
            normalization_type: 归一化类型，可选'minmax'、'zscore'或'pct_change'
            figsize: 图表大小
            save_path: 保存路径，如果为None则不保存
            dpi: 图像分辨率
        """
        # 获取无衰减权重处理的数据
        no_decay_data = self.process_data_no_decay()
        
        # 选择对应的归一化数据列
        if normalization_type == 'minmax':
            price_col = 'close_price_norm_minmax'
            sentiment_col = 'sentiment_norm_minmax'
            type_desc = "MinMax归一化"
            ylabel = "归一化值 (0-1)"
        elif normalization_type == 'zscore':
            price_col = 'close_price_norm_zscore'
            sentiment_col = 'sentiment_norm_zscore'
            type_desc = "Z-score标准化"
            ylabel = "标准化值 (z-score)"
        elif normalization_type == 'pct_change':
            price_col = 'close_price_pct_change'
            sentiment_col = 'sentiment_pct_change'
            type_desc = "相对变化率"
            ylabel = "变化率 (%)"
        else:
            print(f"不支持的归一化类型: {normalization_type}")
            return
        
        # 创建图表
        fig, axes = plt.subplots(2, 2, figsize=figsize)
        
        # 1. 走势图对比 - 有衰减权重
        ax1 = axes[0, 0]
        decay_data = self.merged_data.dropna(subset=[price_col, sentiment_col])
        
        ax1.plot(decay_data.index, decay_data[price_col], 'r-', label='股价', linewidth=2)
        ax1.plot(decay_data.index, decay_data[sentiment_col], 'b-', label='情感指数(带衰减)', linewidth=2)
        
        # 设置x轴日期格式
        ax1.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
        plt.setp(ax1.xaxis.get_majorticklabels(), rotation=45)
        
        # 添加网格线和图例
        ax1.grid(True, linestyle='--', alpha=0.6)
        ax1.legend(loc='best')
        
        # 计算相关系数
        decay_corr = decay_data[price_col].corr(decay_data[sentiment_col])
        ax1.set_title(f'带时间衰减权重 - 相关系数: {decay_corr:.4f}')
        ax1.set_xlabel('日期')
        ax1.set_ylabel(ylabel)
        
        # 2. 走势图对比 - 无衰减权重
        ax2 = axes[0, 1]
        no_decay_data_clean = no_decay_data.dropna(subset=[price_col, sentiment_col])
        
        ax2.plot(no_decay_data_clean.index, no_decay_data_clean[price_col], 'r-', label='股价', linewidth=2)
        ax2.plot(no_decay_data_clean.index, no_decay_data_clean[sentiment_col], 'g-', label='情感指数(无衰减)', linewidth=2)
        
        # 设置x轴日期格式
        ax2.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
        plt.setp(ax2.xaxis.get_majorticklabels(), rotation=45)
        
        # 添加网格线和图例
        ax2.grid(True, linestyle='--', alpha=0.6)
        ax2.legend(loc='best')
        
        # 计算相关系数
        no_decay_corr = no_decay_data_clean[price_col].corr(no_decay_data_clean[sentiment_col])
        ax2.set_title(f'无时间衰减权重 - 相关系数: {no_decay_corr:.4f}')
        ax2.set_xlabel('日期')
        ax2.set_ylabel(ylabel)
        
        # 3. 滞后分析 - 有衰减权重
        ax3 = axes[1, 0]
        max_lag = 5
        
        # 计算不同滞后期的相关系数 - 有衰减
        decay_correlations = []
        for i in range(max_lag + 1):
            if i == 0:
                corr = decay_data[price_col].corr(decay_data[sentiment_col])
            else:
                corr = decay_data[price_col].corr(decay_data[sentiment_col].shift(-i))
            decay_correlations.append(corr)
        
        ax3.bar(range(max_lag + 1), decay_correlations, color='blue', alpha=0.7)
        ax3.plot(range(max_lag + 1), decay_correlations, 'ro-', linewidth=2)
        
        ax3.set_xlabel('滞后天数')
        ax3.set_ylabel('相关系数')
        ax3.set_title('时间衰减权重 - 滞后相关性分析')
        ax3.set_xticks(range(max_lag + 1))
        ax3.axhline(y=0, color='k', linestyle='--', alpha=0.3)
        
        # 添加数值标签
        for i, v in enumerate(decay_correlations):
            ax3.text(i, v + 0.02 * (1 if v >= 0 else -1), f'{v:.3f}', ha='center')
        
        ax3.grid(True, linestyle='--', alpha=0.6)
        
        # 4. 滞后分析 - 无衰减权重
        ax4 = axes[1, 1]
        
        # 计算不同滞后期的相关系数 - 无衰减
        no_decay_correlations = []
        for i in range(max_lag + 1):
            if i == 0:
                corr = no_decay_data_clean[price_col].corr(no_decay_data_clean[sentiment_col])
            else:
                corr = no_decay_data_clean[price_col].corr(no_decay_data_clean[sentiment_col].shift(-i))
            no_decay_correlations.append(corr)
        
        ax4.bar(range(max_lag + 1), no_decay_correlations, color='green', alpha=0.7)
        ax4.plot(range(max_lag + 1), no_decay_correlations, 'ro-', linewidth=2)
        
        ax4.set_xlabel('滞后天数')
        ax4.set_ylabel('相关系数')
        ax4.set_title('无时间衰减权重 - 滞后相关性分析')
        ax4.set_xticks(range(max_lag + 1))
        ax4.axhline(y=0, color='k', linestyle='--', alpha=0.3)
        
        # 添加数值标签
        for i, v in enumerate(no_decay_correlations):
            ax4.text(i, v + 0.02 * (1 if v >= 0 else -1), f'{v:.3f}', ha='center')
        
        ax4.grid(True, linestyle='--', alpha=0.6)
        
        # 调整布局
        plt.suptitle(f'时间衰减权重 vs 无衰减权重对比分析 ({type_desc})', fontsize=16)
        plt.tight_layout()
        plt.subplots_adjust(top=0.93)
        
        # 保存图表
        if save_path:
            plt.savefig(save_path, dpi=dpi, bbox_inches='tight')
            print(f"衰减权重对比分析图已保存至: {save_path}")
        
        plt.show()
        
        # 找出最大相关系数对应的滞后期
        decay_max_corr_idx = np.argmax(np.abs(decay_correlations))
        decay_max_corr = decay_correlations[decay_max_corr_idx]
        
        no_decay_max_corr_idx = np.argmax(np.abs(no_decay_correlations))
        no_decay_max_corr = no_decay_correlations[no_decay_max_corr_idx]
        
        print(f"带时间衰减权重: 最强相关性在滞后 {decay_max_corr_idx} 天时出现，相关系数为 {decay_max_corr:.4f}")
        print(f"无时间衰减权重: 最强相关性在滞后 {no_decay_max_corr_idx} 天时出现，相关系数为 {no_decay_max_corr:.4f}")
        
        return {
            'with_decay': {
                'correlation': decay_corr,
                'lag_correlations': decay_correlations,
                'max_lag': decay_max_corr_idx,
                'max_correlation': decay_max_corr
            },
            'no_decay': {
                'correlation': no_decay_corr,
                'lag_correlations': no_decay_correlations,
                'max_lag': no_decay_max_corr_idx,
                'max_correlation': no_decay_max_corr
            }
        }

    def plot_trend_no_decay(self, title="股价与情感指数走势图(无时间衰减权重)", figsize=(14, 8), save_path=None, dpi=100):
        """
        绘制不带时间衰减权重的股价和情感指数走势图
        
        参数:
            title: 图表标题
            figsize: 图表大小
            save_path: 保存路径，如果为None则不保存
            dpi: 图像分辨率
        """
        # 获取无衰减权重处理的数据
        no_decay_data = self.process_data_no_decay()
        
        if no_decay_data is None or len(no_decay_data) == 0:
            print("没有可用的数据进行绘图")
            return
        
        # 创建图表
        fig, ax1 = plt.subplots(figsize=figsize)
        
        # 绘制股价曲线
        color = 'tab:red'
        ax1.set_xlabel('日期')
        ax1.set_ylabel('股价', color=color)
        ax1.plot(no_decay_data.index, no_decay_data['close_price'], color=color, linewidth=2, marker='o', label='股价')
        ax1.tick_params(axis='y', labelcolor=color)
        
        # 设置x轴日期格式
        ax1.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d'))
        plt.xticks(rotation=45)
        
        # 创建第二个y轴
        ax2 = ax1.twinx()
        color = 'tab:green'
        ax2.set_ylabel('情感指数(无衰减)', color=color)
        ax2.plot(no_decay_data.index, no_decay_data['sentiment_index'], color=color, linewidth=2, marker='x', label='情感指数(无衰减)')
        ax2.tick_params(axis='y', labelcolor=color)
        
        # 添加网格线
        ax1.grid(True, linestyle='--', alpha=0.6)
        
        # 添加图例
        lines1, labels1 = ax1.get_legend_handles_labels()
        lines2, labels2 = ax2.get_legend_handles_labels()
        ax1.legend(lines1 + lines2, labels1 + labels2, loc='upper left')
        
        # 计算相关系数
        corr = no_decay_data['close_price'].corr(no_decay_data['sentiment_index'])
        ax2.annotate(f"相关系数: {corr:.4f}", xy=(0.02, 0.95), xycoords='axes fraction',
                   fontsize=10, bbox=dict(boxstyle="round,pad=0.3", fc="white", ec="gray", alpha=0.8))
        
        # 设置标题
        plt.title(title, fontsize=15)
        
        # 调整布局
        plt.tight_layout()
        
        # 保存图表
        if save_path:
            plt.savefig(save_path, dpi=dpi, bbox_inches='tight')
            print(f"无衰减权重走势图已保存至: {save_path}")
        
        # 显示图表
        plt.show()
        
        return corr
    
    def plot_correlation_no_decay(self, figsize=(14, 10), save_path=None, dpi=100):
        """
        绘制不带时间衰减权重的股价和情感指数相关性分析图
        
        参数:
            figsize: 图表大小
            save_path: 保存路径，如果为None则不保存
            dpi: 图像分辨率
        """
        # 获取无衰减权重处理的数据
        no_decay_data = self.process_data_no_decay()
        
        if no_decay_data is None or len(no_decay_data) < 5:
            print("数据不足，无法进行相关性分析")
            return
        
        # 创建图表
        fig, axes = plt.subplots(2, 2, figsize=figsize)
        
        # 1. 计算相关系数
        correlation = no_decay_data['close_price'].corr(no_decay_data['sentiment_index'])
        corr_text = f"相关系数: {correlation:.4f}"
        
        # 2. 散点图和回归线
        sns.regplot(
            x='sentiment_index', 
            y='close_price', 
            data=no_decay_data, 
            ax=axes[0, 0], 
            scatter_kws={'alpha':0.6, 's':50},
            line_kws={'color': 'red'}
        )
        axes[0, 0].set_title(f'股价与情感指数散点图 ({corr_text})')
        axes[0, 0].set_xlabel('情感指数(无衰减)')
        axes[0, 0].set_ylabel('股价')
        
        # 3. 热力图
        sns.heatmap(
            no_decay_data[['close_price', 'sentiment_index']].corr(),
            annot=True,
            cmap='coolwarm',
            ax=axes[0, 1],
            vmin=-1, vmax=1
        )
        axes[0, 1].set_title('相关性热力图(无衰减)')
        
        # 4. 股价变化与情感指数的关系
        # 计算每日股价变化百分比 (已在process_data_no_decay中计算)
        
        # 去除NaN值
        data_change = no_decay_data.dropna(subset=['price_change_pct', 'sentiment_pct_change'])
        
        # 散点图和回归线
        change_corr = data_change['price_change_pct'].corr(data_change['sentiment_pct_change'])
        change_corr_text = f"相关系数: {change_corr:.4f}"
        
        sns.regplot(
            x='sentiment_pct_change', 
            y='price_change_pct', 
            data=data_change, 
            ax=axes[1, 0], 
            scatter_kws={'alpha':0.6, 's':50},
            line_kws={'color': 'red'}
        )
        axes[1, 0].set_title(f'股价变化百分比与情感指数变化 ({change_corr_text})')
        axes[1, 0].set_xlabel('情感指数变化百分比 (%)')
        axes[1, 0].set_ylabel('股价变化百分比 (%)')
        
        # 5. 情感指数分布直方图
        sns.histplot(no_decay_data['sentiment_index'], kde=True, ax=axes[1, 1])
        axes[1, 1].set_title('情感指数分布(无衰减)')
        axes[1, 1].set_xlabel('情感指数')
        axes[1, 1].set_ylabel('频率')
        
        # 调整布局
        plt.suptitle("股价与情感指数相关性分析(无时间衰减权重)", fontsize=16)
        plt.tight_layout()
        plt.subplots_adjust(top=0.93)
        
        # 保存图表
        if save_path:
            plt.savefig(save_path, dpi=dpi, bbox_inches='tight')
            print(f"无衰减权重相关性分析图已保存至: {save_path}")
        
        # 显示图表
        plt.show()
        
        return {
            'direct': correlation,
            'change': change_corr
        }

def main():
    """主函数"""
    # 文件路径
    stock_data_path = "../data/stock_data/hk01810_stock_data(3).csv"
    sentiment_data_path = "../data/sentiment_index2.csv"
    
    # 创建可视化器，使用时间衰减因子0.9（可调整）
    visualizer = StockSentimentVisualizer(stock_data_path, sentiment_data_path, decay_factor=0.9)
    
    # 原始走势图
    visualizer.plot_trend(title="小米集团(01810)股价与情感指数走势图（带时间衰减权重）",
                          save_path="result2/stock_sentiment_trend_weighted.png")
    
    # 绘制归一化后的走势图
    print("\n=== 归一化走势图分析 ===")
    for norm_type in ['minmax', 'zscore', 'pct_change']:
        corr = visualizer.plot_trend_normalized(
            normalization_type=norm_type,
            save_path=f"result2/stock_sentiment_trend_{norm_type}.png"
        )
        print(f"{norm_type}归一化后的相关系数: {corr:.4f}")
    
    # 绘制归一化后的相关性分析图
    print("\n=== 归一化相关性分析 ===")
    corr_results = visualizer.plot_correlation_normalized(
        save_path="result2/stock_sentiment_correlation_normalized.png"
    )
    for norm_type, corr in corr_results.items():
        print(f"{norm_type}归一化方法的相关系数: {corr:.4f}")
    
    # 原始相关性分析
    visualizer.plot_correlation(save_path="result2/stock_sentiment_correlation_weighted.png")
    
    # 进行滞后分析
    print("\n=== 原始数据滞后分析 ===")
    visualizer.lag_analysis(max_lag=5, save_path="result2/stock_sentiment_lag_analysis_weighted.png")
    
    # 对归一化数据进行滞后分析
    print("\n=== 归一化数据滞后分析 ===")
    for norm_type in ['minmax', 'zscore', 'pct_change']:
        visualizer.lag_analysis_normalized(
            normalization_type=norm_type,
            max_lag=5,
            save_path=f"result2/stock_sentiment_lag_analysis_{norm_type}.png"
        )
    
    # 比较不同归一化方法下的滞后效应
    print("\n=== 归一化方法滞后效应比较 ===")
    visualizer.compare_normalizations_lag_effect(
        max_lag=5,
        save_path="result2/normalization_lag_comparison.png"
    )
    
    # 比较有/无时间衰减权重的分析效果
    print("\n=== 时间衰减权重 vs 无衰减权重对比分析 ===")
    for norm_type in ['minmax', 'zscore', 'pct_change']:
        print(f"\n使用 {norm_type} 归一化方法:")
        results = visualizer.compare_decay_vs_no_decay(
            normalization_type=norm_type,
            save_path=f"result2/decay_vs_no_decay_{norm_type}.png"
        )
        
        # 输出两种方法的对比结果
        decay_results = results['with_decay']
        no_decay_results = results['no_decay']
        
        print(f"带时间衰减权重 - 当前相关系数: {decay_results['correlation']:.4f}, 最大相关系数: {decay_results['max_correlation']:.4f} (滞后 {decay_results['max_lag']} 天)")
        print(f"无时间衰减权重 - 当前相关系数: {no_decay_results['correlation']:.4f}, 最大相关系数: {no_decay_results['max_correlation']:.4f} (滞后 {no_decay_results['max_lag']} 天)")
    
    visualizer.time_scale_analysis(save_path="result2/time_scale_analysis_weighted.png")
    visualizer.cross_correlation_analysis()
    # visualizer.granger_causality_test()
    print(visualizer.extreme_sentiment_event_analysis().head())
    
    # 可视化时间衰减权重
    visualizer.visualize_time_decay(days=10, save_path="result2/time_decay_weights.png")
    
    # 对比不同衰减因子的效果
    plt.figure(figsize=(14, 8))
    for i, factor in enumerate([0.5, 0.7, 0.9, 0.95]):
        plt.subplot(2, 2, i+1)
        dates = pd.date_range(end=pd.Timestamp.now().date(), periods=10)
        days_diff = [(dates[-1] - date).days for date in dates]
        weights = np.power(factor, days_diff)
        weights = weights / weights.sum()
        plt.bar(range(10), weights, alpha=0.7)
        plt.plot(range(10), weights, 'ro-', linewidth=2)
        plt.title(f'衰减因子 = {factor}')
        plt.xlabel('距离交易日的天数')
        plt.ylabel('权重')
        plt.xticks(range(10), [f'T-{d}' for d in days_diff])
        plt.grid(True, linestyle='--', alpha=0.4)
    
    plt.tight_layout()
    plt.savefig("result2/decay_factor_comparison.png", dpi=100)
    plt.show()
    
    # 绘制不带时间衰减权重的分析图
    print("\n=== 不带时间衰减权重的分析 ===")
    
    # 原始走势图（无衰减）
    visualizer.plot_trend_no_decay(title="小米集团(01810)股价与情感指数走势图（无时间衰减权重）",
                               save_path="result2/stock_sentiment_trend_no_decay.png")
    
    # 相关性分析（无衰减）
    corr_results = visualizer.plot_correlation_no_decay(
        save_path="result2/stock_sentiment_correlation_no_decay.png"
    )
    print(f"无衰减权重 - 直接相关系数: {corr_results['direct']:.4f}, 变化率相关系数: {corr_results['change']:.4f}")
    
    # 滞后分析（无衰减，使用Z-score标准化）
    max_lag = 5
    print("\n无衰减权重的滞后分析(Z-score标准化):")
    
    # 准备数据
    no_decay_data = visualizer.process_data_no_decay()
    no_decay_data_clean = no_decay_data.dropna(subset=['close_price_norm_zscore', 'sentiment_norm_zscore'])
    
    # 计算滞后相关性
    correlations = []
    for i in range(max_lag + 1):
        if i == 0:
            corr = no_decay_data_clean['close_price_norm_zscore'].corr(no_decay_data_clean['sentiment_norm_zscore'])
        else:
            corr = no_decay_data_clean['close_price_norm_zscore'].corr(no_decay_data_clean['sentiment_norm_zscore'].shift(-i))
        correlations.append(corr)
        print(f"滞后 {i} 天: 相关系数 = {corr:.4f}")
    
    # 找出最大相关系数
    max_idx = np.argmax(np.abs(correlations))
    print(f"最大相关性在滞后 {max_idx} 天，相关系数为 {correlations[max_idx]:.4f}")

if __name__ == "__main__":
    main() 